# -*-coding:utf-8 -*-
#!/usr/bin/env python
#function: 图像预处理
#input:真彩png图像
#output:预处理后的二值化图像

import cv2
import numpy as np
from math import atan
from scipy import ndimage

    

#倾斜校正，输入二值化图像,输出校正后二值图像
def slant_correction(src):
    #膨胀
    kel = np.ones(40)
    dilate = cv2.dilate(src,kel)
    
    #轮廓
    image, contours, hier = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    #cv2.imshow('dilate',dilate)


    #寻找最大轮廓
    area=0
    cnt = contours[0]
    for c in contours:
        tmp_area = cv2.contourArea(c)
        if tmp_area>=area:
            area = tmp_area
            cnt = c

    #利用最小外界矩形，提取倾斜角
    rect = cv2.minAreaRect(cnt)

    
    #图像倾斜校正
    add_angle = -1.305244 #补偿角
    if rect[1][0]<rect[1][1] and rect[2]!=-90:
        angle = rect[2]+add_angle 
    else:
        angle = rect[2]+90+add_angle

    print('angle:%f'%angle)
    img = ndimage.rotate(src, angle)
    #cv2.imshow("slant_cor",img)
    return img

def pre_process(src):
    gray = cv2.cvtColor(src,cv2.COLOR_BGR2GRAY) #灰度化
    (T,img) = cv2.threshold(gray, 230, 255, cv2.THRESH_BINARY)#阈值化
    img = cv2.medianBlur(img,3)  #中值滤波
    #img = cv2.bilateralFilter(img,3,20,20) #双边滤波
    
    img = slant_correction(img) #倾斜校正
    
    #形态学闭运行
    img = cv2.dilate(img,(3,3))
    img = cv2.erode(img,(3,3))

    (T,img) = cv2.threshold(img, 200, 255, cv2.THRESH_BINARY)#再次阈值化

    #膨胀
    kenel = np.ones(2)
    img = cv2.dilate(img,kenel)

    
    return img

